import os
from typing import TypedDict, List

from langgraph.constants import START, END
from langgraph.graph import StateGraph

from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class TopicState(TypedDict):
    query: str
    queries: List[str]
    documents: List[str]
    keywords: List[str]
    sentiments: List[str]
    summaries: List[str]
    report: str

class TopicAgent:

    def __init__(self, llm_model, agent_tools):
        self.llm_model = llm_model
        self.agent_tools = agent_tools

    def research_topic(self, state: TopicState):
        logger.info("TopicAgent research_topic start")
        template = """
            将这个研究主题分解为3-5个具体的、 可搜索的查询：
            {topic}
    
            让每个查询都集中且可操作。
            以中文回答。
            
            请按照以下格式输出：
            1.xxxxx
            
            2.xxxxx
            
            3.xxxxx
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        response = chain.invoke({ "topic": state["query"] })
        logger.info(f"TopicAgent research_topic response: \n{response}")
        return { "queries": [query for query in response.split("\n")] }

    def search_document(self, state: TopicState):
        logger.info("TopicAgent search_document start")
        documents = []
        queries = state["queries"]
        for query in queries:
            for tool in self.agent_tools:
                if "search_web" == tool.name:
                    result = tool.func(query)
                    documents.append(result)
        return { "documents": documents }

    def extract_keyword(self, state: TopicState):
        logger.info("TopicAgent extract_keyword start")
        template = """
            提取以下文本的关键词：
            {document}
            只返回提取的关键字，不带任何说明
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        documents = state["documents"]
        keywords = []
        for document in documents:
            response = chain.invoke({ "document": document })
            keywords.append(response)
        logger.info(f"TopicAgent extract_keyword keywords: {keywords}")
        return { "keywords": keywords }

    def analyze_sentiment(self, state: TopicState):
        logger.info("TopicAgent analyze_sentiment start")
        template = """
            分析以下文本的情感：
            {document}

            以中文回答情感是正面还是负面，不带任何说明。
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        documents = state["documents"]
        sentiments = []
        for document in documents:
            response = chain.invoke({ "document": document })
            sentiments.append(response)
        logger.info(f"TopicAgent analyze_sentiment sentiments: {sentiments}")
        return { "sentiments": sentiments }

    def summary_document(self, state: TopicState):
        logger.info("TopicAgent summary_document start")
        template = """
            总结以下文本的内容：
            {document}
            只返回总结的结果，不带任何说明。
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        documents = state["documents"]
        summaries = []
        for document in documents:
            response = chain.invoke({ "document": document })
            summaries.append(response)
        return { "summaries": summaries }

    def generate_report(self, state: TopicState):
        logger.info("TopicAgent generate_report start")
        keywords = state["keywords"]
        sentiments = state["sentiments"]
        summaries = state["summaries"]
        contents = []
        for idx, (keyword, sentiment, summary) in enumerate(zip(keywords, sentiments, summaries)):
            contents.append(f"{idx}. 关键字：{keyword}，情感：{sentiment}，总结：{summary}\n\n")
        context = "\n".join(contents)
        logger.info(f"TopicAgent generate_report context len: {len(context)}")

        template = """
            你是一位主题报告撰写专家。请根据主题、关键字、情感和总结，撰写一篇严谨详实的主题报告。
            
            主题：{topic}
            关键字、情感和总结: {context}
            
            请用中文详尽撰写。
            
            以markdown格式输出：
            一、封面（Cover Page）
            报告标题（清晰、准确、简洁，可含副标题）
            作者姓名
            所属单位（学校、公司、机构）
            提交对象（如：指导老师、部门领导、客户名称）
            完成日期（年月日）
            示例：
            《人工智能在医疗影像诊断中的应用现状与挑战》
            作者：
            单位：XX大学计算机学院
            提交至：国家自然科学基金项目组
            日期：2025年11月20日
                
            二、摘要（Abstract）（200–500字）
            目的：研究要解决什么问题？
            方法：采用什么研究方法或技术路线？
            结果：主要发现或数据是什么？
            结论：得出什么结论？有何意义或建议？
            要求：独立成文，不引用图表/参考文献，避免缩写（首次出现需全称）。
                
            三、关键词（Keywords）（3–8个）
            选取能代表报告核心内容的专业术语
            中文报告常用中文关键词，英文报告用英文
            示例：人工智能；医学影像；深度学习；辅助诊断；模型可解释性
            
            四、目录（Table of Contents）
            自动生成（推荐使用 Word / LaTeX 的目录功能）
            包含章节标题及对应页码
            可包含图表目录（如报告中图表较多）
            
            五、引言 / 绪论（Introduction）
            研究背景与意义（为什么值得研究？）
            国内外研究现状综述（简要回顾相关工作，指出空白）
            研究目标与核心问题
            报告结构安排（“本文共分X章，第一章……”）
            目的：让读者理解“为什么要研究这个问题”以及“本报告如何展开”。
            
            六、研究方法 / 理论框架（Methodology / Theoretical Framework）
            研究设计（定性/定量/混合？案例研究/实验/调查？）
            数据来源（问卷、数据库、实验采集、公开数据集等）
            分析工具或模型（如 Python、SPSS、Qdrant、BERT 等）
            技术路线图（可选，用流程图展示步骤）
            伦理说明（如涉及人类受试者）
            要求：可复现性——他人应能根据描述重复你的研究。
                
            七、研究结果 / 分析与讨论（Results / Analysis and Discussion）
            客观呈现结果（用文字、表格、图表）
            对结果进行解释和分析（为什么会出现这样的结果？）
            与已有研究对比（支持还是矛盾？原因是什么？）
            讨论结果的合理性、局限性、意外发现
            建议：
            图表需有编号和标题（如“图1：准确率对比”）
            避免在结果部分直接下结论，留到“结论”章节
            
            八、结论与建议（Conclusion and Recommendations）
            总结主要研究发现（呼应引言中的问题）
            强调研究的理论或实践价值
            提出具体建议（政策建议、技术改进、管理措施等）
            指出研究局限性
            展望未来研究方向
            避免：引入新数据或新观点！
                
            九、参考文献（References）
            按照统一格式列出所有引用文献
            常见格式：
            学术论文：APA、MLA、Chicago、GB/T 7714（中国国标）
            工程/技术报告：IEEE 格式
            要求：真实、完整、规范
            示例（GB/T 7714）：
            [1] 刘知远, 孙茂松. 大模型时代的自然语言处理[J]. 中文信息学报, 2023, 37(6): 1–15.
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        response = chain.invoke({ "topic": state["query"], "context": context })
        logger.info(f"TopicAgent generate_report response len: {len(response)}")
        return { "report": response }

    def build_graph(self):
        logger.info("TopicAgent build_graph start")
        graph = StateGraph(TopicState)
        graph.add_node("research_topic", self.research_topic)
        graph.add_node("search_document", self.search_document)
        graph.add_node("extract_keyword", self.extract_keyword)
        graph.add_node("analyze_sentiment", self.analyze_sentiment)
        graph.add_node("summary_document", self.summary_document)
        graph.add_node("generate_report", self.generate_report)

        graph.add_edge(START, "research_topic")
        graph.add_edge("research_topic", "search_document")
        graph.add_edge("search_document", "extract_keyword")
        graph.add_edge("search_document", "analyze_sentiment")
        graph.add_edge("search_document", "summary_document")
        graph.add_edge("extract_keyword", "generate_report")
        graph.add_edge("analyze_sentiment", "generate_report")
        graph.add_edge("summary_document", "generate_report")
        graph.add_edge("generate_report", END)

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agent_TopicReport_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"TopicAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        report = response.get("report", None)
        logger.info(f"TopicAgent invoke report len: {len(report)}")
        return { "chain_result": report }
